from typing import List, Dict

from langchain_core.documents import Document

from library.common import constants
from langchain_openai.chat_models import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from library.modelio.core import get_prompt

chat_model = ChatOpenAI()
embedding_model = OpenAIEmbeddings()


def get_response(chat_history: List[Dict[str, str]], docs: List[Document], query: str):
    """
    @Params: chat_history 沟通历史
    @Params: docs 检索文档
    @Params: query 用户查询
    """
    prompt = get_prompt(chat_history, docs, query)
    return chat_model.invoke(prompt).content


if __name__ == '__main__':
    history = [
        {
            'role': 'user',
            'content': '你叫什么名字'
        },
        {
            'role': 'assistant',
            'content': '长江'
        },
        {
            'role': 'user',
            'content': '我叫李黄河'
        },
        {
            'role': 'assistant',
            'content': '李黄河，您好'
        }
    ]
    docs = [
        Document('123'),
        Document('456'),
        Document('789'),
    ]
    query = '我叫什么名字'
    response = chat_model.invoke(get_prompt(history, docs, query)).content
    print(response)
